import { Annotation, AnnotationRoot, StateGraph } from "@langchain/langgraph";
import { IWorkflowNode, IWorkflowNodeEndEventArgs } from "./IWorkflowNode";
import { SendMessageData } from "./sendMessageData";
import { IWorkflow } from "./IWorkflow";
import { StartNode } from "./baseNodes/startNode";
import { formatDuration, GenID } from "../utils";
import { BaseWorkflowInputVars } from "./baseWorkflowInputVars";
import { DBService } from "../services/dbService";
import { IWorkflowLog } from "../services/domain/IWorkflowLog";
import { EndNode } from "./baseNodes/endNode";

export abstract class WorkflowBase implements IWorkflow {
    protected convID: string;
    abstract WorkflowID: string;
    private logSvr = DBService.NodeLogService;
    protected nodes = new Map<string, IWorkflowNode>();
    protected OnNodeMessage(sender: IWorkflowNode, args: SendMessageData) {
        const msg = `data:${JSON.stringify(args)}\n\n`;
        this.OnLLmMessage && this.OnLLmMessage(msg);
    };
    protected workflow: StateGraph<typeof this.WorkflowVars.State>;
    OnLLmMessage: (chunk: any) => void;
    OnLLmEnd: (args: any) => void;
    abstract WorkflowVars: AnnotationRoot<typeof BaseWorkflowInputVars.spec>;
    protected abstract InitNode(): void;
    protected abstract InitEdge(): void;
    protected _StartNode = new StartNode();
    protected _EndNode = new EndNode();
    protected AddNode(node: IWorkflowNode) {
        if (this.nodes.has(node.name)) {
            throw new Error(`Node ${node.name} already exists`);
        }
        node.OnMessage = this.OnNodeMessage.bind(this);
        node.workflowID = this.WorkflowID;
        if (!node.outputVars || !node.outputVars.State) {
            node.outputVars = Annotation.Root({
                [node.ResultField]: Annotation<string>,
            });
            node.outputVars.State = {
                [node.ResultField]: ""
            };
        }

        this.nodes.set(node.name, node);
    }

    /**
     * 添加节点连线
     *
     * @protected
     * @param {(IWorkflowNode | "__start__")} from 开始节点
     * @param {(IWorkflowNode | "__end__")} to 结束节点
     * @memberof WorkflowBase
     */
    protected AddEdge(...nodes: (IWorkflowNode | "__start__" | "__end__")[]) {
        if (nodes.length < 2) {
            throw new Error("At least two nodes are required to create an edge");
        }

        for (let i = 0; i < nodes.length - 1; i++) {
            let from = nodes[i] as any;
            let to = nodes[i + 1] as any;

            if (typeof from !== "string") {
                from = from.name;
            }

            if (typeof to !== "string") {
                to = to.name;
            }

            console.log(`添加连线：(${from})---->(${to})`);
            this.workflow.addEdge(from as any, to as any);
        }
    }

    /**
     * 添加判断节点
     *
     * @protected
     * @template T
     * @param {IWorkflowNode} startNode 开始节点
     * @param {string} sourceVarkey 需要判断的变量名称
     * @param {(sourceValue: T) => Promise<string>} conditionHandler 条件判断函数
     * @memberof WorkflowBase
     */
    protected AddConditionEdge<T>(
        startNode: IWorkflowNode,
        sourceVarkey: string,
        conditionHandler: (sourceValue: T) => Promise<IWorkflowNode>
    ) {
        this.workflow.addConditionalEdges(startNode.name as any, async (state) => {
            const result = await conditionHandler(state[sourceVarkey]);
            return result.name;
        })
    }

    private beforeNodeRun(node: IWorkflowNode) {
        return {
            timeStart: Date.now(),
        }
    }

    private async afterNodeRun(node: IWorkflowNode, timeStart: number, error?) {
        let logItem: IWorkflowLog = {
            id: GenID(10),
            nodeName: node.name,
            userID: node.inputVars.State.userID,
            convID: node.inputVars.State.convID,
            runID: "",
            workflowID: this.WorkflowID,
            completion_tokens: node.usage.completion_tokens,
            prompt_tokens: node.usage.prompt_tokens,
            total_tokens: node.usage.total_tokens,
            input: node.inputVars ? JSON.stringify(node.inputVars.State) : "",
            output: node.outputVars ? JSON.stringify(node.outputVars.State) : "",
            time: formatDuration(Date.now() - timeStart),
            error,
            createdAt: new Date(),
        };

        console.log(`\n -- ${logItem.workflowID} -> ${logItem.nodeName}(${logItem.time}): convID:${logItem.convID} total:${logItem.total_tokens}`);

        // await this.logSvr.AddLog(logItem);
    }

    protected DefineGraph() {
        console.log("初始化：", this.WorkflowID);
        this.AddNode(this._StartNode);
        this.AddNode(this._EndNode);
        this.InitNode();
        let inputStateDef = {};
        this.nodes.forEach(node => {
            Object.assign(inputStateDef, node.inputVars.spec);
            Object.assign(inputStateDef, node.outputVars.spec);
        })
        this.WorkflowVars = Annotation.Root(inputStateDef as any);
        this.workflow = new StateGraph(this.WorkflowVars);

        this.nodes.forEach(node => {
            console.log("添加node：", node.name);
            this.workflow.addNode(node.name, async (state: typeof this.WorkflowVars.State) => {
                node.inputVars.State = { ...state };
                let beforeResult = this.beforeNodeRun(node);
                try {
                    const result = await node.Run();
                    this.afterNodeRun(node, beforeResult.timeStart);
                    return result;
                } catch (error) {
                    this.afterNodeRun(node, beforeResult.timeStart, error);
                    throw error;
                }
            });
        });

        this.AddEdge("__start__", this._StartNode);
        this.InitEdge();

        return this.workflow.compile();
    }

    protected OnRunEnd() {
    }

    public async Run(inputState: typeof this.WorkflowVars.State) {
        this.convID = inputState.convID;
        const app = await this.DefineGraph();

        await this.DrawMermaid(app);

        const resultState = await app.invoke(inputState);
        this.OnRunEnd();
        return resultState;
    }

    Stop() {
        this.nodes.forEach(n => {
            n.Stop();
        });
    }

    async DrawMermaid(app) {
        const graph = await app.getGraphAsync();
        const code = graph.drawMermaid();
        console.log(code);
    }
}